/*=========================================================================
 *
 *  Copyright Arash Akbarinia
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

/**
* @file TrainWithOpenCv.cpp
*
*/

#include "io.h"
#include "laperconfig.h"

#include <opencv2/core/core.hpp>
#include <opencv2/features2d/features2d.hpp>
#include <opencv2/nonfree/features2d.hpp>
#include <opencv2/opencv.hpp>

cv::DenseFeatureDetector detector ( 1, 1, 1, 1, 0 );
int nfeatures = 64;

/// brief
int bytes = 16;
cv::BriefDescriptorExtractor BriefExtractor ( bytes );

/// brisk
int thresh = 30;
int octaves = 3;
float pattern_scale = 1.0f;
cv::BRISK brisk_extractor ( thresh, octaves, pattern_scale );

/// freak
bool orientation_normalised = true;
bool scale_normalised = true;
float freak_pattern_scale = 22.0f;
int noctaves = 4;
std::vector< int > selected_pairs;
cv::FREAK freak_extractor ( orientation_normalised, scale_normalised, freak_pattern_scale, noctaves, selected_pairs );

/// orb
float scale_factor = 1.2f;
int orb_nlevels = 8;
int orb_edge_threshold = 31;
int first_level = 0;
int wta_k = 2;
int score_type = cv::ORB::HARRIS_SCORE;
int patch_size = 3;
cv::OrbDescriptorExtractor orb_extractor ( nfeatures, scale_factor, orb_nlevels, orb_edge_threshold, first_level, wta_k, score_type, patch_size );

/// sift
int noctave_layers = 3;
double contrast_threshold = 0.04;
double edge_threshold = 10;
double sigma = 1.6;
cv::SiftDescriptorExtractor sift_extractor ( nfeatures, noctave_layers, contrast_threshold, edge_threshold, sigma );

/// surf
double hessian_threshold = 100;
int surf_noctaves = 4;
int surf_noctave_layers = 3;
bool extended = true;
bool upright = false;
cv::SurfDescriptorExtractor surf_extractor ( hessian_threshold, surf_noctaves, surf_noctave_layers, extended, upright );

void ExctractFeaturesSift ( const cv::Mat& image, cv::Mat& feature )
{
  // getting the features
  std::vector< cv::KeyPoint > keypoints;
  detector.detect ( image, keypoints );

  sift_extractor.compute ( image, keypoints, feature );

  feature.convertTo ( feature, CV_32FC1 );
}

void ExctractFeaturesSurf ( const cv::Mat& image, cv::Mat& feature )
{
  // getting the features
  std::vector< cv::KeyPoint > keypoints;
  detector.detect ( image, keypoints );

  surf_extractor.compute ( image, keypoints, feature );

  feature.convertTo ( feature, CV_32FC1 );
}

void ExctractFeaturesBrisk ( const cv::Mat& image, cv::Mat& feature )
{
  // getting the features
  std::vector< cv::KeyPoint > keypoints;
  detector.detect ( image, keypoints );

  brisk_extractor.compute ( image, keypoints, feature );

  feature.convertTo ( feature, CV_32FC1 );
}

void ExctractFeaturesBrief ( const cv::Mat& image, cv::Mat& feature )
{
  // getting the features
  std::vector< cv::KeyPoint > keypoints;
  detector.detect ( image, keypoints );

  // NOTE: BriefDescriptorExtractor only accepts CV_8U
  cv::Mat UcharImage;
  image.convertTo ( UcharImage, CV_8UC1 );
  BriefExtractor.compute ( UcharImage, keypoints, feature );

  feature.convertTo ( feature, CV_32FC1 );
}

void ExctractFeaturesBriskFreak ( const cv::Mat& image, cv::Mat& feature )
{
  // getting the features
  std::vector< cv::KeyPoint > keypoints;
  detector.detect ( image, keypoints );

  freak_extractor.compute ( image, keypoints, feature );

  feature.convertTo ( feature, CV_32FC1 );
}

void ExtractFeaturesCustom ( const cv::Mat& image, const int z, std::vector < std::vector< float > >& FeatureVector )
{
  cv::Mat SobelX;
  cv::Sobel ( image, SobelX, -1, 1, 0 );
  cv::Mat SobelY;
  cv::Sobel ( image, SobelY, -1, 0, 1 );

  FeatureVector.resize ( image.rows * image.cols );
  for ( int i = 0; i < image.rows; ++i )
    {
      const uint16_t* RowI = image.ptr< uint16_t > ( i );
      const uint16_t* RowISobelX = SobelX.ptr< uint16_t > ( i );
      const uint16_t* RowISobelY = SobelY.ptr< uint16_t > ( i );
      for ( int j = 0; j < image.cols; ++j )
        {
          unsigned int k = ( i * image.cols ) + j;
          uint16_t PixelValue = *RowI++;
          uint16_t PixelValueSobelX = *RowISobelX++;
          uint16_t PixelValueSobelY = *RowISobelY++;

          std::vector< float > FeatureValue ( 6 );
          FeatureValue[0] = i;
          FeatureValue[1] = j;
          FeatureValue[2] = z;
          FeatureValue[3] = PixelValue;
          FeatureValue[4] = PixelValueSobelX;
          FeatureValue[5] = PixelValueSobelY;

          FeatureVector[k] = FeatureValue;
        }
    }
}

void ConvertGtToLabels ( const cv::Mat& gt, std::vector< int >& labels )
{
  labels.resize ( gt.rows * gt.cols );
  for ( int i = 0; i < gt.rows; ++i )
    {
      const uint8_t* RowI = gt.ptr< uint8_t > ( i );
      for ( int j = 0; j < gt.cols; ++j )
        {
          unsigned int k = ( i * gt.cols ) + j;
          uint8_t PixelValue = *RowI++;
          if ( PixelValue == 0 )
            {
              labels[k] = 0;
            }
          else
            {
              labels[k] = 1;
            }
        }
    }
}

void ConvertFeatureMatToVector ( const cv::Mat& FeatureMat, std::vector < std::vector< float > >& FeatureVector )
{
  FeatureVector.resize ( FeatureMat.rows );
  for ( int i = 0; i < FeatureMat.rows; ++i )
    {
      const uint8_t* RowI = FeatureMat.ptr< uint8_t > ( i );
      FeatureVector[i] = std::vector< float > ( RowI, RowI + FeatureMat.cols );
//       for ( int j = 0; j < FeatureMat.cols; ++j )
//         {
// 
//         }
    }
}

void ExtractFeatures ( const std::vector< std::string > ImageFileNames, const std::vector< std::string > GtFileNames, std::vector< int >& labels, std::vector< std::vector< float > >& features, const std::string FeatureType )
{
  unsigned int nimages = ImageFileNames.size();

  if ( nimages == 0 )
    {
      FILE_LOG ( logERROR ) << "provided folder for images is empty";
      exit ( EXIT_FAILURE );
    }
  if ( GtFileNames.size() == 0 )
    {
      FILE_LOG ( logERROR ) << "provided folder for ground truth is empty";
      exit ( EXIT_FAILURE );
    }

  void ( * tocall ) ( const cv::Mat&, cv::Mat& );

  if ( FeatureType.compare ( "brief" ) == 0 )
    {
      tocall = ExctractFeaturesBrief;
    }
  else if ( FeatureType.compare ( "brisk" ) == 0 )
    {
      tocall = ExctractFeaturesBrisk;
    }
  else if ( FeatureType.compare ( "freak" ) == 0 )
    {
      tocall = ExctractFeaturesBriskFreak;
    }
  else if ( FeatureType.compare ( "sift" ) == 0 )
    {
      tocall = ExctractFeaturesSift;
    }
  else if ( FeatureType.compare ( "surf" ) == 0 )
    {
      tocall = ExctractFeaturesSurf;
    }
  else
    {
      FILE_LOG ( logERROR ) << "feature type is not supported";
      exit ( EXIT_FAILURE );
    }

  for ( unsigned int i = 0; i < nimages; i++ )
    {
      FILE_LOG ( logINFO ) << "reading " << ImageFileNames[i];
      cv::Mat image = cv::imread ( ImageFileNames[i], cv::IMREAD_ANYDEPTH );
      cv::Mat gt = cv::imread ( GtFileNames[i], cv::IMREAD_GRAYSCALE );

      if ( image.empty ( ) )
        {
          FILE_LOG ( logERROR ) << "error! couldn't read " << ImageFileNames[i];
          exit ( EXIT_FAILURE );
        }
      if ( gt.empty ( ) )
        {
          FILE_LOG ( logERROR ) << "error! couldn't read " << GtFileNames[i];
          exit ( EXIT_FAILURE );
        }

      cv::Mat FeatureMat;
//       tocall ( image, FeatureMat );

      // adding it to the training data set
      std::vector < std::vector< float > > FeaturesI;
//       ConvertFeatureMatToVector ( FeatureMat, FeaturesI );
      ExtractFeaturesCustom( image, i, FeaturesI );
      features.insert ( features.end(), FeaturesI.begin(), FeaturesI.end() );

      // putting the lable
      std::vector< int > LabelsI;
      ConvertGtToLabels ( gt, LabelsI );
      labels.insert ( labels.end(), LabelsI.begin(), LabelsI.end() );

      // writing the descriptor in the file
      {
        std::string DescriptorFilename = "DescriptorFilename";
//         descriptor_filename.append ( type );
        DescriptorFilename.append ( FeatureType );
//         descriptor_filename.append ( stri );
        DescriptorFilename.append ( ".yaml" );
        cv::FileStorage fs ( DescriptorFilename, cv::FileStorage::WRITE );
        fs << "PixelDescriptors" << "[";
        for ( unsigned int j = 0; j < FeaturesI.size(); j++ )
          {
            fs << "{:" << "label" << LabelsI[j] << "descriptor" << FeaturesI[j] << "}";
          }
        fs.release ( );
      }
    }
}

cv::Mat real_all_descriptors ( int argc, char** argv )
{
  std::vector< cv::Mat > all_descriptors;
  int ncolumns = 0;
  for ( int i = 5; i < argc; i++ )
    {
      cv::FileStorage fs;
      fs.open ( argv[i], cv::FileStorage::READ );
      cv::Mat tmpdescriptors;
      fs["descriptors"] >> tmpdescriptors;
      ncolumns += tmpdescriptors.cols;
      all_descriptors.push_back ( tmpdescriptors );
      fs.release();
//       cout << "read image " << argv[i] << " size " << all_descriptors.back().rows << " x " << all_descriptors.back().cols << endl;
    }

//   cout << "ncolumns " << ncolumns << endl;
  cv::Mat descriptors ( atoi ( argv[4] ) * 2, ncolumns, CV_32FC1 );
  int copycolumn = 0;
  for ( unsigned int i = 0; i < all_descriptors.size(); i++ )
    {
//       cout << "copycolumn " << copycolumn << endl;
      all_descriptors[i].copyTo ( cv::Mat ( descriptors, cv::Rect ( copycolumn, 0, all_descriptors[i].cols, all_descriptors[i].rows ) ) );
      copycolumn += all_descriptors[i].cols;
    }
  return descriptors;
}

int main ( int argc, char** argv )
{

  if ( argc < 5 )
    {
      FILE_LOG ( logERROR ) << "usage " << argv[0] << ": InputImageDirectory InputGroundTruthDirectory OutputDirectory FeatureType ";
      return EXIT_FAILURE;
    }

  std::string InputImageDirectory = argv[1];
  std::string InputGtDirectory = argv[2];
  std::string OutputDirectory = argv[3];
  // TODO: check fi the feature is accepted
  std::string FeatureType = argv[4];

  std::vector< std::string > ImageFileNames;
  std::vector< std::string > GtFileNames;
  std::vector< std::string > FileExtensions;
  FileExtensions.push_back ( ".png" );
  laper::io::ReadFilesOfDirectory ( InputImageDirectory, ImageFileNames, FileExtensions, true );
  laper::io::ReadFilesOfDirectory ( InputGtDirectory, GtFileNames, FileExtensions, true );

  /// extracting features
  std::vector< std::vector< float > > features;
  std::vector< int > labels;
  ExtractFeatures ( ImageFileNames, GtFileNames, labels, features, FeatureType );
  std::cout << features.size() << " " << labels.size() << std::endl;

  return 0;
}
